import heapq


class Solution:
    def trapRainWater(self, heightMap) -> int:
        next = [(0, -1), (0, 1), (-1, 0), (1, 0)]
        dp = [[-1 for y in range(len(heightMap[x]))] for x in range(len(heightMap))]
        changeHeap = []

        for i in range(len(heightMap[0])):
            heapq.heappush(changeHeap, (heightMap[0][i], 0, i))
            heapq.heappush(changeHeap, (heightMap[len(heightMap) - 1][i], len(heightMap) - 1, i))
            dp[0][i] = heightMap[0][i]
            dp[len(heightMap) - 1][i] = heightMap[len(heightMap) - 1][i]
        for i in range(1, len(heightMap) - 1):
            heapq.heappush(changeHeap, (heightMap[i][0], i, 0))
            heapq.heappush(changeHeap, (heightMap[i][len(heightMap[i]) - 1], i, len(heightMap[i]) - 1))
            dp[i][0] = heightMap[i][0]
            dp[i][len(heightMap[i]) - 1] = heightMap[i][len(heightMap[i]) - 1]

        while changeHeap:
            item = heapq.heappop(changeHeap)
            for t in next:
                i = item[1] + t[0]
                j = item[2] + t[1]
                if 1 <= i < len(heightMap)-1 and 1 <= j < len(heightMap[i])-1:
                    if dp[i][j] == -1:
                        newHeight = max(item[0], heightMap[i][j])
                    else:
                        newHeight = max(min(item[0], dp[i][j]), heightMap[i][j])
                    if newHeight != dp[i][j]:
                        dp[i][j] = newHeight
                        heapq.heappush(changeHeap, (newHeight, i, j))
        r = 0
        for i in range(1, len(heightMap) - 1):
            for j in range(1, len(heightMap[i]) - 1):
                r += dp[i][j] - heightMap[i][j]

        return r


if __name__ == '__main__':
    so = Solution()
    num = [
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

    r = so.trapRainWater(num)
    print(r)
